#!/bin/bash
startfunc(){
    mkdir -p docs
	mkdir -p docs/imgs
	mkdir -p ../origin_data
	mkdir -p ../data
	mkdir -p ../logs
	mkdir -p ../backup
}


lintfunc(){
    autopep8 --in-place *.py
    autopep8 --in-place */*.py
}

clearfunc(){
    rm -rf __pycache__
    rm -rf */__pycache__
    rm -rf */*/__pycache__
}

cleanfunc(){
    clearfunc
    rm -rf docs/*.log
    rm -rf docs/*.md
    rm -rf docs/imgs/*.png
	rm -rf ../logs/*
}

backupfunc(){
    mkdir -p ../backup
    mkdir -p ../backup/logs
    mkdir -p ../backup/docs
    mkdir -p ../backup/docs/imgs
    cp ../logs/* ../backup/logs
    cp docs/*.md ../backup/docs
    cp docs/imgs/*.png ../backup/docs/imgs
}

cmd=$1
shift 1
case $cmd in
    start)
        startfunc
        ;;
    lint)
        lintfunc
        ;;
    clear)
        clearfunc
        ;;
    clean)
        cleanfunc
        ;;
    doc)
        taskid=Null
        while getopts "t:w:" opt; do
            case $opt in
                t)
                    taskid=$OPTARG
                    ;;
            esac
        done
        python3 gendoc.py -t $taskid
        ;;
    docs)
        task_file=docs/tasklists.log
        cat $task_file | xargs -I {} python3 gendoc.py -t {}
        ;;
    backup)
        backupfunc
        ;;
    run)
        log_file="docs/cmd_history.log"
        gpuid=0
        while getopts "t:m:g:" opt; do
            case $opt in
                t)
                    # get param after -t, example: tasks=1,2,3
                    tasks=$OPTARG
                    # seprate use ',', example: task_list=(1,2,3)
                    task_list=(${tasks//,/ })
                    ;;
                m)
                    comment=$OPTARG
                    ;;
                g)
                    gpuid=$OPTARG
                    ;;
            esac
        done
        # record this cmd
        now=`date +%F_%T`
        echo "[$now]: [./run.sh -t $tasks -g $gpuid] [$comment]" >> $log_file
        # Displays upcoming commands
        echo "> TRAINING TASKS: [$tasks] will be exec one-by-one on gpu:$gpuid"

        for tsk in ${task_list[@]}; do
            python3 main.py -i $tsk -g $gpuid
            sleep 1m
        done
        ;;
    prepare)
        datasets='cifar_10,stl_10'
        while getopts "d:" opt; do
            case $opt in
                d)
                    datasets=$OPTARG
                    ;;
            esac
        done
        data_list=(${datasets//,/ })
        for data in ${data_list[@]}; do
            python3 prepare.py -d $data
        done
        ;;
    model)
        cp ../backup/logs/*.pth conf/
        cp ../backup/logs/*.pt conf/
        ;;
    unmodel)
        rm -f conf/*.pth
        rm -f conf/*.pt
        ;;
esac
